#!/usr/bin/python3
# -*- coding: utf-8 -*-

"""
# @version: v1.1
# @author : cd
# @Email : 19688513@qq.com
# @Project : new-horizons-engine
# @File : StockBasicInfoManager.py
# @Software: PyCharm
# @time: 2025/5/30 9:51
# @description : 股票基本信息管理类
"""

import urllib3
import logging
import sqlite3
from datetime import datetime
import pandas as pd
from tqdm import tqdm

from AKShareDataFetcher import AKShareDataFetcher
from DatabaseManager import DatabaseManager

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
# 禁用不安全的请求警告
urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)


class StockBasicInfoManager:
    """股票基本信息管理类"""

    def __init__(self, db_manager: DatabaseManager):
        self.db_manager = db_manager
        self.fetcher = AKShareDataFetcher()  # 初始化数据获取器

    def upsert_basic_info(self, stock_code: str, stock_name: str, esg_info: dict = None) -> bool:
        """更新或插入股票基本信息 - 与DatabaseManager表结构保持一致"""
        try:
            # 获取股票详细信息
            info_df = self.fetcher.fetch_stock_detail_info(stock_code)

            if info_df.empty:
                logging.warning(f"无法获取股票 {stock_code} 的详细信息")
                return False

            # 转换为字典格式
            info_dict = info_df.set_index('item')['value'].to_dict()

            # 处理ESG信息
            rating_agency = "秩鼎" if esg_info else "无"
            rating = esg_info.get('ESG评分', "无") if esg_info else "无"
            rating_quarter = esg_info.get('评分日期', '未评级') if esg_info else '未评级'

            # 构造插入数据
            insert_data = (
                stock_code,
                stock_name,
                self._convert_to_float(info_dict.get('总市值', 0)),
                self._convert_to_float(info_dict.get('流通市值', 0)),
                info_dict.get('行业', '未知'),
                info_dict.get('上市时间', '1970-01-01'),
                self._convert_to_int(info_dict.get('总股本', '0')),
                self._convert_to_int(info_dict.get('流通股', '0')),
                rating_agency,
                rating,
                rating_quarter,
                self.get_stock_exchange(stock_code),
                datetime.now().strftime("%Y-%m-%d %H:%M:%S")
            )

            # 使用事务执行插入或更新
            with self.db_manager.get_connection() as conn:
                conn.execute('''
                    INSERT INTO stock_basic_info VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                    ON CONFLICT(stock_code) DO UPDATE SET
                        stock_name = excluded.stock_name,
                        total_market_value = excluded.total_market_value,
                        circulating_market_value = excluded.circulating_market_value,
                        industry = excluded.industry,
                        listing_date = excluded.listing_date,
                        total_shares = excluded.total_shares,
                        circulating_shares = excluded.circulating_shares,
                        rating_agency = excluded.rating_agency,
                        rating = excluded.rating,
                        rating_quarter = excluded.rating_quarter,
                        market_type = excluded.market_type,
                        last_updated = excluded.last_updated
                ''', insert_data)
                conn.commit()

            logging.info(f"股票{stock_code}基本信息更新成功")
            return True
        except Exception as e:
            logging.error(f"股票{stock_code}基本信息更新失败: {str(e)}", exc_info=True)
            return False

    def _convert_to_int(self, value: str) -> int:
        """将字符串值转换为整数（处理带单位的数值）"""
        try:
            if isinstance(value, (int, float)):
                return int(value)

            # 处理带单位的数值
            if '亿' in value:
                num = float(value.replace('亿', '').replace(',', ''))
                return int(num * 100000000)
            elif '万' in value:
                num = float(value.replace('万', '').replace(',', ''))
                return int(num * 10000)
            return int(float(value.replace(',', '')))
        except Exception:
            return 0

    def _convert_to_float(self, value: str) -> float:
        """将字符串值转换为浮点数（处理带单位的数值）"""
        try:
            if isinstance(value, (int, float)):
                return float(value)

            # 处理带单位的数值
            if '亿' in value:
                num = float(value.replace('亿', '').replace(',', ''))
                return num * 100000000
            elif '万' in value:
                num = float(value.replace('万', '').replace(',', ''))
                return num * 10000
            return float(value.replace(',', ''))
        except Exception:
            return 0.0

    def get_stock_exchange(self, code: str) -> str:
        """根据股票代码判断所属交易所"""
        code_str = str(code).zfill(6)
        if code_str.startswith(('6', '9')):
            return 'sh'
        elif code_str.startswith(('0', '3')):
            return 'sz'
        elif code_str.startswith('4') or code_str.startswith('8'):
            return 'bj'
        else:
            return 'unknown'

    def get_stock_codes(self, market_types: list = None, excluded_prefixes: list = None) -> pd.DataFrame:
        """获取符合条件的股票代码"""
        if market_types is None:
            market_types = ['sz', 'sh']
        if excluded_prefixes is None:
            excluded_prefixes = ['688', '300']

        # 构建SQL查询条件
        market_condition = "market_type IN ({})".format(", ".join(f"'{mt}'" for mt in market_types))
        exclusion_conditions = " AND ".join(f"stock_code NOT LIKE '{prefix}%'" for prefix in excluded_prefixes)

        query = f'''
            SELECT DISTINCT stock_code ,listing_date
            FROM stock_basic_info
            WHERE {market_condition} 
            AND {exclusion_conditions};
        '''

        try:
            with self.db_manager.get_connection() as conn:
                cursor = conn.execute(query)
                results = cursor.fetchall()
                return pd.DataFrame(results, columns=['stock_code', 'listing_date'])
        except sqlite3.OperationalError as e:
            logging.warning(f"查询股票代码失败: {e}")
            return pd.DataFrame(columns=['stock_code', 'listing_date'])

    def initialize_stock_database(self):
        """初始化股票数据库"""
        logging.info("开始初始化股票数据库...")

        # 获取所有股票基本信息
        logging.info("获取股票基本信息...")
        basic_info = self.fetcher.fetch_stock_basic_info(max_retries=3, retry_delay=2, request_delay=0.4)

        # 获取ESG信息
        logging.info("获取ESG评级信息...")
        esg_info_df  = self.fetcher.fetch_stock_esg_info(max_retries=3, retry_delay=3, request_delay=0.5)
        # 创建代码到ESG信息的映射
        esg_data_dict = {}
        # 去掉 ESG 数据中的后缀
        # 创建代码到ESG信息的映射
        esg_data_dict = {}

        # 去掉 ESG 数据中的后缀并创建映射
        for _, row in esg_info_df.iterrows():
            clean_code = row['股票代码'].replace('.SH', '').replace('.SZ', '').replace('.HK', '')
            esg_data_dict[clean_code] = {
                'ESG评分': row['ESG评分'],
                'E评分': row['环境总评'],
                'S评分': row['社会责任总评'],
                'G评分': row['治理总评'],
                '评分日期': row['评分日期']
            }

                                                                                                               # 转换数据格式
        stock_list = basic_info.to_dict('records')
        total_count = len(stock_list)

        logging.info(f"开始处理{total_count}只股票...")
        success_count = 0
        failed_stocks = []

        # 使用tqdm显示进度条
        for stock in tqdm(stock_list, desc="🚀 更新股票信息", unit="股"):
            stock_code = stock['code']
            stock_name = stock['name']
            # 获取该股票的ESG信息
            esg_data = esg_data_dict.get(stock_code, None)

            if self.upsert_basic_info(stock_code, stock_name, esg_data):
                success_count += 1
            else:
                failed_stocks.append(stock_code)

        # 打印结果统计
        logging.info(f"股票数据库初始化完成! 成功更新{success_count}/{total_count}只股票信息")
        if failed_stocks:
            logging.warning(
                f"以下{len(failed_stocks)}只股票更新失败: {', '.join(failed_stocks[:10])}{'...' if len(failed_stocks) > 10 else ''}")


if __name__ == "__main__":
    # 创建数据库管理器实例
    db_manager = DatabaseManager()

    # 创建股票基本信息管理实例
    stock_manager = StockBasicInfoManager(db_manager)

    # 初始化股票数据库
    stock_manager.initialize_stock_database()

    # 示例：获取符合条件的股票代码
    logging.info("获取符合条件的股票代码...")
    market_types = ['sz', 'sh']
    excluded_prefixes = ['688', '300']
    filtered_stock_codes = stock_manager.get_stock_codes(market_types, excluded_prefixes)

    logging.info(f"符合条件的股票数量: {len(filtered_stock_codes)}")
    logging.info(f"示例股票代码: {filtered_stock_codes.head(5).to_string(index=False)}")

    # 示例：更新单个股票信息
    # test_stock_code = "000001"
    # logging.info(f"更新测试股票 {test_stock_code} 信息...")
    # stock_manager.upsert_basic_info(
    #     stock_code=test_stock_code,
    #     stock_name="平安银行",
    #     esg_info={'ESG评分': 85.2, '评分日期': '2025-05-30'}
    # )